import shutil
import os
import os.path as osp

num_dict = {}
sets = ['train', 'val']
for image_set in sets:
    # 创建图像分类文件夹
    image_set_browser = 'D:\games_project\yolov5-master-f\MYdata\\images_cls' + image_set  # 字符串拼接
    if os.path.exists(image_set_browser):
        shutil.rmtree(image_set_browser)  # 递归删除user\tester 目录的内容
        print('Deleted previous {} file and created a new one'.format(image_set_browser))
    os.makedirs(image_set_browser)

    # 创建标注分类文件夹
    json_set_browser = 'D:\games_project\yolov5-master-f\MYdata\\labels_cls\\' + image_set
    if os.path.exists(json_set_browser):
        shutil.rmtree(json_set_browser)
        print('Deleted previous {} file and created a new one'.format(json_set_browser))
    os.makedirs(json_set_browser)
    #读取类别txt文件夹下的 train val
    image_ids = open('E:\\Test2\\MYdata\\txt\\{}.txt'.format(image_set)).read().strip().split()

    print(image_set)
    for i, image_id in enumerate(image_ids):
        print(i + 1, image_id)
        img = 'E:\\Test2\\all_pic\\{}.jpg'.format(image_id)
        json = 'E:\\Test2\\all_text\\{}.txt'.format(image_id)
        shutil.copy(img, image_set_browser)
        shutil.copy(json, json_set_browser)
    print('img set: {} num: {}'.format(image_set, len(os.listdir(image_set_browser))))
    print('json set: {} num: {} \n'.format(image_set, len(os.listdir(json_set_browser))))
    num_dict[image_set] = len(os.listdir(json_set_browser))  # 创建字典
print("image and json files classify Done")
for key, value in num_dict.items():  # 利用字典，进行显示
    print('{} num: {}'.format(key, value))